import json
from urllib import parse
from urllib.parse import parse_qs

from asgiref.sync import sync_to_async
from django.conf import settings
from django.contrib.auth import get_user_model
from jwt import decode as jwt_decode
from django.db import close_old_connections
from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
from rest_framework_simplejwt.tokens import UntypedToken
from django.urls import resolve
from .connection import WebSocket


def websockets(app):
    async def asgi(scope, receive, send):
        if scope["type"] == "websocket":
            match = resolve(scope["raw_path"])
            new_scope, e = await verify_token(scope)
            if e:
                await send({"type": "websocket.accept"})
                await send({"type": "websocket.send", "text": json.dumps({"code": 401, "msg": "Token is invalid or expire"})})
                return
            await match.func(WebSocket(new_scope, receive, send), *match.args, **match.kwargs)
            return
        await app(scope, receive, send)

    return asgi


async def verify_token(scope):
    # Get the token
    token = parse_qs(scope["query_string"].decode("utf8"))["token"][0]
    print('token:', token)

    # Then token is valid, decode it
    try:
        # This will automatically validate the token and raise an error if token is invalid
        UntypedToken(token)
    except (InvalidToken, TokenError) as e:
        # Token is invalid
        print(e)
        return None, e
    else:
        #  Then token is valid, decode it
        decoded_data = jwt_decode(token, settings.SECRET_KEY, algorithms=["HS256"])
        print("token payload:", decoded_data)
        # Will return a dictionary like -
        # {
        #     "token_type": "access",
        #     "exp": 1568770772,
        #     "jti": "5c15e80d65b04c20ad34d77b6703251b",
        #     "user_id": 6
        # }
        # Get the user using ID
        user = await get_current_user(decoded_data["user_id"])
        print("current user:", user)
        return dict(scope, user=user), None


@sync_to_async
def get_current_user(pk):
    return get_user_model().objects.get(pk=pk)


class TokenAuthMiddleware:
    """
    Custom token auth middleware
    """

    def __init__(self, inner):
        # Store the ASGI application we were passed
        self.inner = inner

    def __call__(self, scope):

        # Close old database connections to prevent usage of timed out connections
        close_old_connections()

        # Get the token
        token = parse_qs(scope["query_string"].decode("utf8"))["token"][0]

        # Try to authenticate the user
        try:
            # This will automatically validate the token and raise an error if token is invalid
            UntypedToken(token)
        except (InvalidToken, TokenError) as e:
            # Token is invalid
            print(e)
            return None
        else:
            #  Then token is valid, decode it
            decoded_data = jwt_decode(token, settings.SECRET_KEY, algorithms=["HS256"])
            print(decoded_data)
            # Will return a dictionary like -
            # {
            #     "token_type": "access",
            #     "exp": 1568770772,
            #     "jti": "5c15e80d65b04c20ad34d77b6703251b",
            #     "user_id": 6
            # }

            # Get the user using ID
            user = get_user_model().objects.get(id=decoded_data["user_id"])

        # Return the inner application directly and let it run everything else
        return self.inner(dict(scope, user=user))
